import pandas as pd
import os
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from tqdm import tqdm
import time

# 从项目模块中导入数据提供函数，以确保一致性
from data_provider import get_market_data as get_project_market_data

# --- 全局配置 ---
TRADES_LOG_PATH = 'all_trades_log.csv'
CACHE_DIR = 'feather_cache'
CHARTS_DIR = 'charts'
MARKET_INDEX_SYMBOL = 'sh.000300' # 使用沪深300指数作为大盘参考

# --- 本地化配置 ---
COLUMN_TRANSLATION = {
    'is_market_ok': '大盘高于MA20',
    'dist_from_high_20d': '距20日高点距离(%)',
    'return_5d': '5日回报率(%)',
    'bias_10_20': '10日与20日均线乖离率(%)',
    'market_down_days_5d': '买后5日大盘下跌天数',
    'max_drawdown_10d': '买后10日内最大跌幅(%)',
    'bounced_then_loss': '曾盈利后亏损',
    'bearish_days_ratio': '持仓期阴线比例(%)',
    'dist_from_ma10_on_buy': '买入时与MA10距离(%)',
    'max_profit_before_loss': '转亏前最大利润(%)',
    'days_in_profit': '盈利持仓天数',
    'market_rsi_on_buy': '买入时大盘RSI',
    'optimal_sell_day_k_shape': '最佳卖点上影线占比',
    'optimal_sell_day_close_vs_ma5': '最佳卖点收盘价高于MA5',
    'optimal_sell_day_volume_ratio': '最佳卖点成交量比率',
    'optimal_sell_day_rsi': '最佳卖点个股RSI'
}

def get_market_index_data(start_date, end_date, adjust_flag):
    """
    包装项目自身的数据获取函数，以适配本脚本的需求。
    注意：这里的日期格式是 'YYYY-MM-DD'
    """
    # 修正调用方式：get_project_market_data 获取的是全量数据
    # 我们在这里进行日期切片
    full_market_df = get_project_market_data(CACHE_DIR, MARKET_INDEX_SYMBOL, adjust_flag)

    if full_market_df is not None and not full_market_df.empty:
        full_market_df['date'] = pd.to_datetime(full_market_df['date'])
        full_market_df.set_index('date', inplace=True)
        # 根据本脚本需要的日期范围进行切片
        return full_market_df.loc[start_date:end_date]
    return None

def calculate_indicators_for_trade(trade_row, stock_df, market_df):
    """为单笔交易计算其买入时的相关技术指标和市场背景，以及买入后的走势特征"""
    buy_date = pd.to_datetime(trade_row['buy_datetime'])
    sell_date = pd.to_datetime(trade_row['sell_datetime'])
    buy_price = trade_row['buy_price']
    
    try:
        actual_date = stock_df.index.asof(buy_date)
        if pd.isna(actual_date): return None
        buy_index_in_stock = stock_df.index.get_loc(actual_date)
    except KeyError:
        return None
    
    if buy_index_in_stock < 60:
        return None
        
    # --- 1. 买入时静态指标 (与之前类似) ---
    stock_slice_before = stock_df.iloc[:buy_index_in_stock + 1]
    
    try:
        actual_market_date = market_df.index.asof(buy_date)
        if pd.isna(actual_market_date):
            is_market_ok = np.nan
        else:
            market_slice = market_df.loc[:actual_market_date]
            market_ma20 = market_slice['close'].rolling(window=20).mean().iloc[-1]
            is_market_ok = market_slice['close'].iloc[-1] > market_ma20
    except (KeyError, IndexError):
        is_market_ok = np.nan

    high_in_20d = stock_slice_before['high'].rolling(window=20).max().iloc[-1]
    dist_from_high_20d = (stock_slice_before['close'].iloc[-1] / high_in_20d - 1) * 100 if high_in_20d > 0 else 0
    return_5d = (stock_slice_before['close'].iloc[-1] / stock_slice_before['close'].iloc[-6] - 1) * 100 if buy_index_in_stock >= 5 else 0
    ma10 = stock_slice_before['close'].rolling(window=10).mean().iloc[-1]
    ma20 = stock_slice_before['close'].rolling(window=20).mean().iloc[-1]
    bias_10_20 = (ma10 / ma20 - 1) * 100 if ma20 > 0 else 0
    # 新增：买入时与10日均线距离
    dist_from_ma10_on_buy = (buy_price / ma10 - 1) * 100 if ma10 > 0 else 0

    # --- 2. 买入后动态指标 (新增核心逻辑) ---
    
    # 2.1 市场背景分析
    market_slice_after = market_df.loc[buy_date:buy_date + pd.Timedelta(days=10)]
    if len(market_slice_after) > 1:
        market_returns = market_slice_after['close'].pct_change().dropna()
        market_down_days_5d = (market_returns.head(5) < 0).sum()
    else:
        market_down_days_5d = np.nan
    
    # 新增：计算买入时大盘RSI
    try:
        actual_market_date = market_df.index.asof(buy_date)
        if pd.isna(actual_market_date):
            market_rsi_on_buy = np.nan
        else:
            market_slice_for_rsi = market_df.loc[:actual_market_date]
            if len(market_slice_for_rsi) >= 15:
                delta = market_slice_for_rsi['close'].diff()
                gain = delta.where(delta > 0, 0)
                loss = -delta.where(delta < 0, 0)
                avg_gain = gain.ewm(span=14, adjust=False).mean().iloc[-1]
                avg_loss = loss.ewm(span=14, adjust=False).mean().iloc[-1]
                if avg_loss > 0:
                    rs = avg_gain / avg_loss
                    market_rsi_on_buy = 100 - (100 / (1 + rs))
                else:
                    market_rsi_on_buy = 100
            else:
                market_rsi_on_buy = np.nan
    except (KeyError, IndexError):
        market_rsi_on_buy = np.nan

    # 2.2 个股买后走势
    stock_slice_after = stock_df.loc[buy_date:sell_date]
    if len(stock_slice_after) > 1:
        # 买后N日内最大跌幅
        ten_day_slice = stock_slice_after.head(10)
        max_drawdown_10d = (ten_day_slice['low'].min() / buy_price - 1) * 100
        
        # 是否曾盈利后亏损 (仅对亏损股有意义)
        bounced_then_loss = 1 if stock_slice_after['high'].max() > buy_price else 0

        # 持仓期阴线比例
        bearish_days = (stock_slice_after['close'] < stock_slice_after['open']).sum()
        bearish_days_ratio = (bearish_days / len(stock_slice_after)) * 100

        # 新增：计算转亏前最大利润 和 盈利持仓天数 (仅对假突破交易有意义)
        max_profit_before_loss = np.nan
        days_in_profit = np.nan
        # 新增：初始化最佳卖点指标
        optimal_sell_day_k_shape = np.nan
        optimal_sell_day_close_vs_ma5 = np.nan
        optimal_sell_day_volume_ratio = np.nan
        optimal_sell_day_rsi = np.nan

        if bounced_then_loss == 1 and trade_row['pnl'] <= 0:
            # 计算最大利润
            max_high_after_buy = stock_slice_after['high'].max()
            max_profit_before_loss = (max_high_after_buy / buy_price - 1) * 100
            
            # 计算盈利天数
            profitable_days = stock_slice_after[stock_slice_after['low'] > buy_price]
            days_in_profit = len(profitable_days)

            # --- 新增：定位最佳卖点并提取特征 ---
            optimal_sell_date = stock_slice_after['high'].idxmax()
            optimal_sell_day_loc = stock_df.index.get_loc(optimal_sell_date)
            
            if optimal_sell_day_loc >= 14: # 确保有足够数据计算指标
                optimal_sell_day_data = stock_df.iloc[optimal_sell_day_loc]
                stock_slice_for_optimal_day = stock_df.iloc[:optimal_sell_day_loc + 1]

                # 1. K线形态特征 (上影线占比)
                k_high = optimal_sell_day_data['high']
                k_low = optimal_sell_day_data['low']
                k_close = optimal_sell_day_data['close']
                k_range = k_high - k_low
                optimal_sell_day_k_shape = (k_high - k_close) / k_range if k_range > 0 else 0

                # 2. 与5日线的关系
                ma5 = stock_slice_for_optimal_day['close'].rolling(5).mean().iloc[-1]
                optimal_sell_day_close_vs_ma5 = 1 if k_close > ma5 else 0

                # 3. 成交量特征
                vol = optimal_sell_day_data['volume']
                avg_vol_5d = stock_slice_for_optimal_day['volume'].shift(1).rolling(5).mean().iloc[-1]
                optimal_sell_day_volume_ratio = vol / avg_vol_5d if avg_vol_5d > 0 else 1

                # 4. 个股RSI
                delta = stock_slice_for_optimal_day['close'].diff()
                gain = delta.where(delta > 0, 0).ewm(span=14, adjust=False).mean()
                loss = -delta.where(delta < 0, 0).ewm(span=14, adjust=False).mean()
                if loss.iloc[-1] > 0:
                    rs = gain.iloc[-1] / loss.iloc[-1]
                    optimal_sell_day_rsi = 100 - (100 / (1 + rs))
                else:
                    optimal_sell_day_rsi = 100
    else:
        max_drawdown_10d = np.nan
        bounced_then_loss = np.nan
        bearish_days_ratio = np.nan
        max_profit_before_loss = np.nan
        days_in_profit = np.nan
        optimal_sell_day_k_shape = np.nan
        optimal_sell_day_close_vs_ma5 = np.nan
        optimal_sell_day_volume_ratio = np.nan
        optimal_sell_day_rsi = np.nan
        
    return {
        # 静态指标
        'is_market_ok': int(is_market_ok) if not np.isnan(is_market_ok) else np.nan,
        'dist_from_high_20d': dist_from_high_20d,
        'return_5d': return_5d,
        'bias_10_20': bias_10_20,
        'dist_from_ma10_on_buy': dist_from_ma10_on_buy,
        'market_rsi_on_buy': market_rsi_on_buy,
        # 动态指标
        'market_down_days_5d': market_down_days_5d,
        'max_drawdown_10d': max_drawdown_10d,
        'bounced_then_loss': bounced_then_loss if trade_row['pnl'] <= 0 else 0, # 此指标只对亏损股统计
        'bearish_days_ratio': bearish_days_ratio,
        'max_profit_before_loss': max_profit_before_loss,
        'days_in_profit': days_in_profit,
        'optimal_sell_day_k_shape': optimal_sell_day_k_shape,
        'optimal_sell_day_close_vs_ma5': optimal_sell_day_close_vs_ma5,
        'optimal_sell_day_volume_ratio': optimal_sell_day_volume_ratio,
        'optimal_sell_day_rsi': optimal_sell_day_rsi,
    }

def plot_distributions(df, columns_to_plot, translations):
    """为指定的列绘制盈利/亏损对比分布图 (中文本地化)"""
    if not os.path.exists(CHARTS_DIR):
        os.makedirs(CHARTS_DIR)
        
    for col in columns_to_plot:
        if col not in df.columns:
            continue
            
        chinese_col_name = translations.get(col, col)

        # 对二元指标(0/1)或计数指标使用条形图更合适
        if df[col].nunique() < 8:
            plt.figure(figsize=(10, 6))
            sns.countplot(data=df, x=col, hue='trade_type', palette={'盈利': 'g', '亏损': 'r'})
        else: # 对连续指标使用直方图
            plt.figure(figsize=(12, 7))
            sns.histplot(data=df, x=col, hue='trade_type', kde=True, bins=40, palette={'盈利': 'g', '亏损': 'r'}, alpha=0.6)
        
        plt.title(f'{chinese_col_name} 在盈利与亏损交易中的分布对比', fontsize=16)
        plt.xlabel(chinese_col_name, fontsize=12)
        plt.ylabel('交易次数', fontsize=12)
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.legend(title='交易类型') # 自动从数据中获取 '盈利' 和 '亏损' 标签
        
        # 文件名也使用中文
        save_path = os.path.join(CHARTS_DIR, f'分布对比_{chinese_col_name}.png')
        plt.savefig(save_path)
        plt.close()
        #print(f"对比图表已保存: {save_path}")

def main():
    """主执行函数"""
    if not os.path.exists(TRADES_LOG_PATH):
        print(f"错误: 交易日志 '{TRADES_LOG_PATH}' 不存在。请先运行 main.py 生成日志。")
        return
        
    trades_df = pd.read_csv(TRADES_LOG_PATH)
    
    # --- 准备数据 ---
    # 扩展日期范围以确保能计算买入后的指标
    min_date = (pd.to_datetime(trades_df['buy_datetime']).min() - pd.Timedelta(days=80)).strftime('%Y-%m-%d')
    max_date = (pd.to_datetime(trades_df['sell_datetime']).max() + pd.Timedelta(days=5)).strftime('%Y-%m-%d')
    
    # 统一从 config.json 读取 adjust_flag (此处假设为 '2')
    market_df = get_market_index_data(min_date, max_date, '2')
    if market_df is None or market_df.empty:
        print("错误: 无法获取大盘指数数据，分析中止。")
        return

    # 添加 'trade_type' 列 (Win/Loss)，并本地化
    trades_df['trade_type'] = np.where(trades_df['pnl'] > 0, '盈利', '亏损')
    
    print(f"共发现 {len(trades_df)} 笔交易。开始回溯情景并计算指标...")
    
    indicator_data_list = []
    
    for _, row in tqdm(trades_df.iterrows(), total=len(trades_df), desc="分析所有交易"):
        symbol = str(row['symbol'])
        
        cache_path = os.path.join(CACHE_DIR, f"{symbol}.feather")
        if not os.path.exists(cache_path):
            continue
            
        stock_df = pd.read_feather(cache_path)
        stock_df.set_index('date', inplace=True)
        
        indicators = calculate_indicators_for_trade(row, stock_df, market_df)
        if indicators:
            indicators['trade_type'] = row['trade_type'] # 把 Win/Loss 标签也加进去
            indicator_data_list.append(indicators)
            
    if not indicator_data_list:
        print("未能为任何交易计算出有效的指标数据。")
        return

    indicator_df = pd.DataFrame(indicator_data_list)
    
    print("\n--- 分组统计摘要 (盈利 vs 亏损) ---")
    
    # 本地化 describe() 的输出
    summary_df = indicator_df.groupby('trade_type').describe()
    stat_translation = {'count': '数量', 'mean': '均值', 'std': '标准差', 'min': '最小值', '25%': '25分位数', '50%': '中位数', '75%': '75分位数', 'max': '最大值'}
    
    summary_df.columns = summary_df.columns.set_levels([COLUMN_TRANSLATION.get(col, col) for col in summary_df.columns.levels[0]], level=0)
    summary_df.columns = summary_df.columns.set_levels([stat_translation.get(stat, stat) for stat in summary_df.columns.levels[1]], level=1)
    summary_df.index.name = '交易类型'
    print(summary_df.T)
    
    # --- 新增：对特定指标进行更深入的文本分析 ---
    loss_df = indicator_df[indicator_df['trade_type'] == '亏损']
    if not loss_df.empty:
        print("\n--- 亏损交易深度分析 ---")
        
        # 1. 分析"假突破"情况 (曾盈利后亏损)
        bounced_count = loss_df['bounced_then_loss'].sum()
        total_losses = len(loss_df)
        bounced_ratio = (bounced_count / total_losses) * 100
        print(f"1. '假突破'分析: 在 {total_losses} 笔亏损交易中，有 {bounced_count} 笔 ({bounced_ratio:.2f}%) 的股价在买入后曾一度高于买入价，但最终仍以亏损告终。")
        
        # 2. 分析市场背景
        avg_market_down_days = loss_df['market_down_days_5d'].mean()
        print(f"2. 市场背景: 亏损交易买入后，大盘在后续5个交易日内平均有 {avg_market_down_days:.2f} 天是下跌的。")

        # 3. 新增：对"假突破"交易进行深度画像
        fake_breakthrough_df = loss_df[loss_df['bounced_then_loss'] == 1]
        if not fake_breakthrough_df.empty:
            print("\n--- '假突破'交易深度画像 ---")
            
            # --- 新增：最佳卖点特征分析 ---
            print("\n  -- 最佳卖点 (利润最大化日) 特征分析 --")
            
            # 连续性指标
            continuous_stats = {
                'optimal_sell_day_k_shape': "上影线占比",
                'optimal_sell_day_volume_ratio': "成交量比率(vs 5日均量)",
                'optimal_sell_day_rsi': "个股RSI"
            }
            for col, name in continuous_stats.items():
                stats = fake_breakthrough_df[col].describe()
                print(f"    - {name}:")
                print(f"      均值: {stats['mean']:.2f}, 中位数: {stats['50%']:.2f}, 75分位数: {stats['75%']:.2f}")

            # 分类指标
            close_vs_ma5_pct = (fake_breakthrough_df['optimal_sell_day_close_vs_ma5'].sum() / len(fake_breakthrough_df.dropna(subset=['optimal_sell_day_close_vs_ma5']))) * 100
            print(f"    - 收盘价高于MA5的比例: {close_vs_ma5_pct:.2f}%")
            
            print("\n  -- 买入时特征回顾 --")
            stats_to_show = {
                'dist_from_ma10_on_buy': "买入时与MA10距离(%)",
                'max_profit_before_loss': "转亏前最大利润(%)",
                'days_in_profit': "盈利持仓天数",
                'market_rsi_on_buy': "买入时大盘RSI"
            }
            for col, name in stats_to_show.items():
                stats = fake_breakthrough_df[col].describe()
                print(f"    - {name}:")
                print(f"      均值: {stats['mean']:.2f}, 中位数: {stats['50%']:.2f}, 标准差: {stats['std']:.2f}, 最大值: {stats['max']:.2f}")

    #print("\n--- 正在生成指标对比分布图 ---")
    columns_to_visualize = list(COLUMN_TRANSLATION.keys())
    plot_distributions(indicator_df, columns_to_visualize, COLUMN_TRANSLATION)
    
    print("\n分析完成！图表已保存到 'charts' 文件夹中。")

if __name__ == '__main__':
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    main() 